import pyautogui
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from matplotlib.colors import to_hex
import matplotlib as mpl

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号


def get_dominant_screen_color(region_size=200, k=3):
    """
    获取屏幕中心区域的主色调
    :param region_size: 中心区域边长（像素）
    :param k: 聚类数量（建议3-5）
    :return: 主色调的十六进制代码和RGB值
    """
    # 获取屏幕尺寸
    screen_width, screen_height = pyautogui.size()

    # 计算中心区域坐标
    center_x, center_y = screen_width // 2, screen_height // 2
    half_size = region_size // 2
    region = (
        center_x - half_size,  # 左上角x
        center_y - half_size,  # 左上角y
        region_size,  # 宽度
        region_size  # 高度
    )

    # 截取屏幕中心区域
    screenshot = pyautogui.screenshot(region=region)
    img_array = np.array(screenshot)

    # 重塑像素数据为2D数组 (每个像素一行，RGB三列)
    pixels = img_array.reshape(-1, 3)

    # 使用K-Means聚类找出主要颜色
    kmeans = KMeans(n_clusters=k, n_init=10, random_state=42)
    kmeans.fit(pixels)

    # 获取聚类中心（主要颜色）
    colors = kmeans.cluster_centers_.astype(int)

    # 获取最大簇的颜色作为主色调
    counts = np.bincount(kmeans.labels_)
    dominant_color = colors[np.argmax(counts)]

    # 转换为十六进制代码
    hex_color = to_hex(dominant_color / 255)

    return hex_color, dominant_color


def visualize_results(dominant_color, region_size=200):
    """可视化结果"""
    # 创建显示区域
    fig, axs = plt.subplots(1, 2, figsize=(10, 5))

    # 显示颜色方块
    color_block = np.ones((100, 100, 3)) * dominant_color / 255
    axs[0].imshow(color_block)
    axs[0].set_title(f'主色调: {to_hex(dominant_color / 255)}')
    axs[0].axis('off')

    # 显示屏幕截图
    try:
        screen_width, screen_height = pyautogui.size()
        center_x, center_y = screen_width // 2, screen_height // 2
        half_size = region_size // 2
        region = (
            center_x - half_size,
            center_y - half_size,
            region_size,
            region_size
        )
        screenshot = pyautogui.screenshot(region=region)
        axs[1].imshow(np.array(screenshot))
        axs[1].set_title('屏幕中心区域')
        axs[1].axis('off')
    except Exception as e:
        print(f"截图显示失败: {e}")

    plt.tight_layout()

    # 保存图像而不是直接显示
    plt.savefig('color_result.png', bbox_inches='tight', dpi=100)
    print("结果已保存为 'color_result.png'")

    # 在PyCharm中安全显示
    try:
        plt.show()
    except Exception as e:
        print(f"显示图表时出错: {e}. 请查看保存的图像文件。")


if __name__ == "__main__":
    # 获取主色调（默认区域200x200像素）
    hex_color, rgb_color = get_dominant_screen_color()

    print(f"主色调 (RGB): {rgb_color}")
    print(f"十六进制代码: {hex_color}")

    # 可视化结果
    visualize_results(rgb_color)